#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
板块股票关联表模型 (BlockInfo)
"""

from sqlalchemy import Column, String, ForeignKey, Index
from sqlalchemy.orm import relationship
from typing import List, Optional, Dict

from .BaseModel import BaseModel

class BlockInfo(BaseModel):
    """板块股票关联表模型"""
    __tablename__ = 'blockinfo'
    
    # 列名映射（CSV列名 -> 数据库字段名）
    aliases = {
        'blockCode': 'block_code',
        'securityId': 'security_id',
        'market': 'market',
        '#板块编号': 'block_code',
        '证券代码': 'security_id',
        '#证券代码': 'security_id',
        '市场': 'market',
        '#市场': 'market'
    }
    
    block_code = Column(String(20), ForeignKey('block.block_code'), primary_key=True)
    security_id = Column(String(20), primary_key=True)
    
    # 关联关系
    block = relationship("Block", back_populates="block_stocks")
    
    # 索引
    __table_args__ = (
        Index('idx_blockinfo_block_code', 'block_code'),
        Index('idx_blockinfo_security_id', 'security_id'),
    )

    @classmethod
    def get_session(cls):
        """获取数据库会话"""
        return cls._db_manager.get_stock_data_session()
    
    @classmethod
    def _get_existing_record(cls, session, data):
        """获取现有记录"""
        return session.query(cls).filter(
            cls.block_code == data['block_code'],
            cls.security_id == data['security_id']
        ).first()
    
    @classmethod
    def _update_existing_record(cls, existing, data):
        """更新现有记录（板块股票关联通常不需要更新）"""
        pass  # 板块股票关联通常不需要更新
    
    @classmethod
    def _to_dict(cls, result):
        """将结果转换为字典"""
        return {
            'block_code': result.block_code,
            'security_id': result.security_id
        }
    
    @classmethod
    def get_stocks_by_block(cls, block_code: str) -> List[str]:
        """根据板块代码获取股票列表"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            securities = session.query(cls.security_id).filter(
                cls.block_code == block_code
            ).all()
            session.close()
            return [sec[0] for sec in securities]
        except Exception as e:
            # 获取板块股票失败: {str(e)}
            session.close()
            return []
    
    @classmethod
    def get_blocks_by_stock(cls, security_id: str) -> List[str]:
        """根据股票代码获取板块列表"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            blocks = session.query(cls.block_code).filter(
                cls.security_id == security_id
            ).all()
            session.close()
            return [block[0] for block in blocks]
        except Exception as e:
            # 获取股票板块失败: {str(e)}
            session.close()
            return []
    
    @classmethod
    def get_industry_by_stock(cls, security_id: str) -> Optional[str]:
        """根据股票代码获取行业名称（category=2的板块）"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            
            # 通过BlockInfo和Block关联查询，找到category=2的板块
            from .Block import Block
            result = session.query(Block.block_name).join(
                cls, cls.block_code == Block.block_code
            ).filter(
                cls.security_id == security_id,
                Block.category == 2
            ).first()
            
            session.close()
            return result[0] if result else None
        except Exception as e:
            # 获取股票行业失败: {str(e)}
            session.close()
            return None
    
    @classmethod
    def get_industries_by_stocks(cls, security_ids: List[str]) -> Dict[str, str]:
        """批量获取股票行业信息"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            
            # 批量查询所有股票的行业信息
            from .Block import Block
            results = session.query(cls.security_id, Block.block_name).join(
                Block, cls.block_code == Block.block_code
            ).filter(
                cls.security_id.in_(security_ids),
                Block.category == 2
            ).all()
            
            session.close()
            
            # 转换为字典格式
            industry_dict = {}
            for security_id, industry_name in results:
                industry_dict[security_id] = industry_name
            
            return industry_dict
        except Exception as e:
            # 批量获取股票行业失败: {str(e)}
            session.close()
            return {}
    
    @classmethod
    def import_from_dataframe(cls, df):
        """从DataFrame导入板块股票关联数据"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            success_count = 0
            
            for _, row in df.iterrows():
                try:
                    # 确保 security_id 为字符串格式，保留前导零
                    security_id = str(row['security_id']).zfill(6)  # 补齐到6位
                    block_code = str(row['block_code'])
                    
                    # 检查是否已存在相同记录
                    existing = session.query(cls).filter(
                        cls.block_code == block_code,
                        cls.security_id == security_id
                    ).first()
                    if not existing:
                        # 创建新记录
                        block_info = cls(
                            block_code=block_code,
                            security_id=security_id
                        )
                        session.add(block_info)
                        success_count += 1
                    
                except Exception as e:
                    # 处理板块股票关联数据失败: {str(e)}
                    continue
            
            session.commit()
            session.close()
            return success_count
            
        except Exception as e:
            # 导入板块股票关联数据失败: {str(e)}
            session.rollback()
            session.close()
            return 0
    
    @classmethod
    def to_dataframe(cls):
        """转换为DataFrame"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            results = session.query(cls).all()
            
            data_list = []
            for result in results:
                data_list.append({
                    'block_code': result.block_code,
                    'security_id': result.security_id
                })
            
            session.close()
            
            import pandas as pd
            df = pd.DataFrame(data_list)
            # 从数据库加载了 {len(df)} 条板块股票关联数据
            return df
            
        except Exception as e:
            # 转换为DataFrame失败: {str(e)}
            session.close()
            import pandas as pd
            return pd.DataFrame()
    
    @classmethod
    def load_from_db(cls, db_path: str = "stock_data.db"):
        """从数据库加载板块股票关联数据（兼容方法）"""
        # 创建临时数据库管理器
        from .BaseModel import DatabaseManager
        db_manager = DatabaseManager(stock_data_db_path=db_path)
        cls.set_db_manager(db_manager)
        
        try:
            df = cls.to_dataframe()
            return df
        finally:
            db_manager.close_all_sessions()
